#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
生成训练收敛曲线对比图

读取所有模型的 training_log.csv 并生成对比图表
"""

import os
import sys
import csv
from pathlib import Path
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')

# 配置中文字体（如果可用）
try:
    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei', 'DejaVu Sans']
    plt.rcParams['axes.unicode_minus'] = False
except:
    pass

BASE_DIR = Path(__file__).parent.parent
MODEL_DIR = BASE_DIR / "model"
RESULTS_DIR = BASE_DIR / "results"
RESULTS_DIR.mkdir(exist_ok=True)


def read_training_log(log_path):
    """读取训练日志"""
    if not log_path.exists():
        return None
    
    steps = []
    losses = []
    
    with open(log_path, 'r', encoding='utf-8') as f:
        reader = csv.DictReader(f)
        for row in reader:
            if 'loss' in row and row['loss']:
                try:
                    steps.append(len(steps) + 1)
                    losses.append(float(row['loss']))
                except:
                    continue
    
    return steps, losses if steps else None


def create_training_comparison():
    """创建训练收敛曲线对比图"""
    
    models = [
        ("Complex LoRA", MODEL_DIR / "text2code_lora_complex" / "training_log.csv", '#1f77b4'),
        ("Multitask LoRA", MODEL_DIR / "multitask_lora" / "training_log.csv", '#ff7f0e'),
        ("Distilled Student", MODEL_DIR / "distilled_student" / "training_log.csv", '#2ca02c'),
        ("Contrastive", MODEL_DIR / "contrastive_pretrained" / "training_log.csv", '#d62728'),
    ]
    
    fig, axes = plt.subplots(2, 2, figsize=(14, 10))
    fig.suptitle('Text2Code Models Training Convergence Comparison', 
                 fontsize=16, fontweight='bold')
    
    # 单独绘制每个模型
    for idx, (name, log_path, color) in enumerate(models):
        ax = axes[idx // 2, idx % 2]
        
        data = read_training_log(log_path)
        
        if data:
            steps, losses = data
            ax.plot(steps, losses, color=color, linewidth=1.5, alpha=0.7)
            ax.set_title(f'{name}', fontsize=12, fontweight='bold')
            ax.set_xlabel('Training Steps', fontsize=10)
            ax.set_ylabel('Loss', fontsize=10)
            ax.grid(True, alpha=0.3)
            
            # 添加最终 loss 标注
            final_loss = losses[-1] if losses else 0
            ax.text(0.95, 0.95, f'Final Loss: {final_loss:.4f}', 
                   transform=ax.transAxes, ha='right', va='top',
                   bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
        else:
            ax.text(0.5, 0.5, 'No Training Log', 
                   transform=ax.transAxes, ha='center', va='center',
                   fontsize=12, color='gray')
            ax.set_title(f'{name}', fontsize=12, fontweight='bold')
    
    plt.tight_layout()
    
    output_path = RESULTS_DIR / "training_convergence_comparison.png"
    plt.savefig(output_path, dpi=150, bbox_inches='tight')
    plt.close()
    
    print(f"[SUCCESS] Training comparison saved: {output_path}")
    
    # 创建综合对比图
    create_combined_comparison(models)


def create_combined_comparison(models):
    """创建综合对比图（所有模型在同一张图上）"""
    
    plt.figure(figsize=(12, 7))
    
    for name, log_path, color in models:
        data = read_training_log(log_path)
        
        if data:
            steps, losses = data
            # 归一化步数到 0-1 范围以便对比
            max_steps = max(steps)
            normalized_steps = [s / max_steps for s in steps]
            plt.plot(normalized_steps, losses, label=name, 
                    color=color, linewidth=2, alpha=0.8)
    
    plt.xlabel('Normalized Training Progress', fontsize=12)
    plt.ylabel('Loss', fontsize=12)
    plt.title('Training Loss Convergence - All Models', 
             fontsize=14, fontweight='bold')
    plt.legend(loc='best', fontsize=10)
    plt.grid(True, alpha=0.3)
    
    output_path = RESULTS_DIR / "training_convergence_combined.png"
    plt.savefig(output_path, dpi=150, bbox_inches='tight')
    plt.close()
    
    print(f"[SUCCESS] Combined comparison saved: {output_path}")


def create_performance_summary():
    """创建性能摘要图表"""
    
    # 数据
    models = ['Complex\nLoRA', 'Multitask\nLoRA', 'Complex\nINT8', 'Multitask\nINT8']
    accuracies = [70, 75, 80, 79]
    memories = [2.0, 2.0, 0.39, 0.39]
    
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))
    
    # 准确率对比
    colors = ['#4CAF50', '#2196F3', '#FF9800', '#9C27B0']
    bars1 = ax1.bar(range(len(models)), accuracies, color=colors, alpha=0.8)
    ax1.set_ylabel('Accuracy (%)', fontsize=12)
    ax1.set_title('Model Accuracy Comparison', fontsize=13, fontweight='bold')
    ax1.set_xticks(range(len(models)))
    ax1.set_xticklabels(models)
    ax1.set_ylim(0, 100)
    ax1.grid(axis='y', alpha=0.3)
    
    # 添加数值标签
    for bar, acc in zip(bars1, accuracies):
        ax1.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 1,
                f'{acc}%', ha='center', va='bottom', fontweight='bold')
    
    # 内存占用对比
    bars2 = ax2.bar(range(len(models)), memories, color=colors, alpha=0.8)
    ax2.set_ylabel('Memory (GB)', fontsize=12)
    ax2.set_title('Memory Usage Comparison', fontsize=13, fontweight='bold')
    ax2.set_xticks(range(len(models)))
    ax2.set_xticklabels(models)
    ax2.set_ylim(0, 2.5)
    ax2.grid(axis='y', alpha=0.3)
    
    # 添加数值标签
    for bar, mem in zip(bars2, memories):
        ax2.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.05,
                f'{mem} GB', ha='center', va='bottom', fontweight='bold')
    
    plt.tight_layout()
    
    output_path = RESULTS_DIR / "performance_summary.png"
    plt.savefig(output_path, dpi=150, bbox_inches='tight')
    plt.close()
    
    print(f"[SUCCESS] Performance summary saved: {output_path}")


def main():
    print("=" * 80)
    print("生成训练收敛曲线对比图")
    print("=" * 80)
    
    create_training_comparison()
    create_performance_summary()
    
    print("\n" + "=" * 80)
    print("[SUCCESS] All charts generated successfully!")
    print("=" * 80)
    print(f"\n查看结果:")
    print(f"  {RESULTS_DIR}")
    print("=" * 80 + "\n")


if __name__ == "__main__":
    main()
